#! -*- coding: utf-8 -*-
"""
@Info: 从标注的excel文件中转换为训练txt数据
"""
import pandas as pd
import json
import schema
from sklearn.model_selection import train_test_split


if __name__ == '__main__':
    df = pd.read_excel('../../data/AIPS语料-20240716.xlsx')
    # df = df[df.iloc[:, 7] == 1 | df.iloc[:, 7] == 2 | df.iloc[:, 7] == 0]  # 筛选已标注用于训练的数据
    # df = df[df.iloc[:, 7] != 10] # 这部分语料不清晰，需确认
    # df.rename(columns={df.columns[2]: 'operate_intent'}, inplace=True)
    df.rename(columns={df.columns[4]: 'operate_intent'}, inplace=True)
    # df.rename(columns={df.columns[4]: 'second_class'}, inplace=True)
    df.rename(columns={df.columns[5]: 'text'}, inplace=True)

    df = df.dropna(subset=['text', 'operate_intent'], how='any')
    df = df.iloc[1:, [5, 4]]

    df_other = pd.read_excel('../../data/语料-other.xlsx')
    df_other.rename(columns={df_other.columns[0]: 'text'}, inplace=True)
    df_other.rename(columns={df_other.columns[1]: 'operate_intent'}, inplace=True)
    df_other = df_other.iloc[1:, :]
    df_other['operate_intent'] = 0

    df = pd.concat([df, df_other])
    df['operate_intent'] = df['operate_intent'].map(str)

    # # 解决可能的数据不均衡问题
    # df_chadan = df[df.loc[:, 'operate_intent'] == '插单']
    # for i in range(10):
    #     df = pd.concat([df, df_chadan])
    # for i in range(2):
    #     df = pd.concat([df, df_other])

    intents = set(df['operate_intent'])
    print('operate_intent:', intents)
    # 查看每个分类的个数
    for intent in intents:
        print(str(intent) + ':', sum(df.loc[:, 'operate_intent'] == intent))

    train_df, eval_df = train_test_split(df, test_size=0.1, random_state=42)
    train_df = df
    
    # 将结果写入文件
    with open('./data/classify_train.txt', 'w', encoding='utf-8') as json_file:
        for _, row in train_df.iterrows():
            row['cats'] = schema.get_schema(row['operate_intent'])
            json_file.write(json.dumps(row.to_dict(), ensure_ascii=False) + '\n')

    with open('./data/classify_eval.txt', 'w', encoding='utf-8') as json_file:
        for _, row in eval_df.iterrows():
            row['cats'] = schema.get_schema(row['operate_intent'])
            json_file.write(json.dumps(row.to_dict(), ensure_ascii=False) + '\n')
